package watcher

import (
    "sync"
    "regexp"
    "os/user"
    "golang.org/x/net/context"
    "fmt"
    "time"
    "strconv"
    "io/ioutil"
    "strings"
    "os/exec"
    "os"
    "bufio"
    "io"
    "encoding/json"
    "../pb"
    "google.golang.org/grpc"
    "github.com/op/go-logging"
)

var log = logging.MustGetLogger("watcher_service")

type ProcWatcherService struct {
    mu              sync.RWMutex
    procs           map[string]*ProcessItem
    RefreshInterval int64
    defineFile      string
    authPort        string
    lastScanTime    int64
    auth            pb.AuthServiceClient
}

func (pws *ProcWatcherService) SetAuth(conn *grpc.ClientConn) {
    if conn != nil {
        log.Infof("setup auth client on: %v", conn)
        pws.auth = pb.NewAuthServiceClient(conn)
    }
}

func NewProcWatcherService(fn string) (*ProcWatcherService, error) {
    f, err := os.Stat(fn)
    if err != nil || f.IsDir() {
        log.Errorf("config file can't open: %v", err)
        return nil, fmt.Errorf("can't open: %v", err)
    }
    b, err := ioutil.ReadFile(fn)
    if err != nil {
        log.Errorf("config file can't open: %v", err)
        return nil, err
    }
    p := &ProcWatcherService{}
    err = json.Unmarshal(b, &p.procs)
    if err != nil {
        return nil, err
    }
    p.defineFile = fn
    p.RefreshInterval = 5
    return p, nil
}

type ProcessItem struct {
    MatchMulti     bool                `json:"match_multi"`
    ProcessPattern string              `json:"pattern"`
    ProcessUser    string              `json:"user"`
    Become         string              `json:"become"`
    Commands       map[string][]string `json:"commands"` //with commands params
    reg            *regexp.Regexp
    process        Process
}

func (pi *ProcessItem) Reset() {
    if pi.MatchMulti {
        pi.process = &MultiProcess{}
    } else {
        pi.process = &SingleProcess{}
    }
}

type Process interface {
    Status() (*pb.Status, error)
    Add(user.User, int64, int) error
    Info(bool) string
    Uptime() (int64)
}

func (pws *ProcWatcherService) Auth(ctx context.Context, token *pb.Token) (bool, error) {
    if pws.auth != nil {
        resp, err := pws.auth.AuthToken(ctx, &pb.AuthTokenRequest{token})
        if err != nil {
            log.Errorf("auth error: %v", err)
            return resp.Valid, err
        }
    }
    return true, nil
}

func (pws *ProcWatcherService) Query(ctx context.Context, req *pb.QueryWatcherRequest) (*pb.QueryWatcherResponse, error) {
    log.Infof("handle query request: %s ,fresh %v", req.ProcName, req.Refresh)
    pws.mu.Lock()
    defer pws.mu.Unlock()
    pass, err := pws.Auth(ctx, req.Token)
    log.Infof("handle query request, token  passed: %v", pass)
    if !pass {
        return nil, fmt.Errorf("Insufficient Privilege")
    }
    proc, ok := pws.procs[req.ProcName]
    if req.Refresh || time.Now().Unix()-pws.lastScanTime > pws.RefreshInterval {
        log.Infof("handle query request, scan start")
        if err := pws.scan(); err != nil {
            return nil, err
        }
    }
    ret := &pb.QueryWatcherResponse{}
    if !ok {
        ret.Err = fmt.Sprintf("proc not found: %s", req.ProcName)
        return ret, nil
    }
    stat, err := proc.process.Status()
    ret.Status = stat
    log.Infof("handle query request finish")
    return ret, err
}

func (pws *ProcWatcherService) List(ctx context.Context, req *pb.ListRequest) (*pb.ListResponse, error) {
    pws.mu.RLock()
    defer pws.mu.RUnlock()
    pass, _ := pws.Auth(ctx, req.Token)
    if !pass {
        return nil, fmt.Errorf("Insufficient Privilege")
    }
    var procNames []string
    for n := range pws.procs {
        procNames = append(procNames, n)
    }
    return &pb.ListResponse{ProcName: procNames}, nil
}

func (pws *ProcWatcherService) Show(ctx context.Context, req *pb.ShowRequest) (*pb.ShowResponse, error) {
    pws.mu.RLock()
    defer pws.mu.RUnlock()
    pass, _ := pws.Auth(ctx, req.Token)
    if !pass {
        return nil, fmt.Errorf("Insufficient Privilege")
    }
    proc, ok := pws.procs[req.ProcName]
    ret := &pb.ShowResponse{Commands: make(map[string]string)}
    if ok {
        for k, v := range proc.Commands {
            ret.Commands[k] = strings.Join(v, " ")
        }
    }
    return ret, nil
}

func (pws *ProcWatcherService) Run(req *pb.RunRequest, srv pb.WatcherService_RunServer) error {
    log.Infof("handle run request from %s, proc name: %s , command name: %s", req.Token.User, req.ProcName, req.CommandName)
    pws.mu.Lock()
    defer pws.mu.Unlock()
    pass, err := pws.Auth(context.Background(), req.Token)
    if !pass {
        return fmt.Errorf("Insufficient Privilege")
    }
    proc, ok := pws.procs[req.ProcName]
    if !ok {
        return fmt.Errorf("no such proc: %s", req.ProcName)
    }
    cmds, ok := proc.Commands[req.CommandName]
    if !ok {
        return fmt.Errorf("no such command: %s", req.CommandName)
    }
    cmd := exec.Command(cmds[0], cmds[1:]...)
    cmd.Env = append(cmd.Env, os.Environ()...)
    stderr, err := cmd.StderrPipe()
    stdout, err := cmd.StdoutPipe()
    output := bufio.NewReader(io.MultiReader(stderr, stdout))
    if err != nil {
        log.Errorf("run command error: %v", err)
        return err
    }
    ts := time.Now().Unix()
    if err := cmd.Start(); err != nil {
        log.Errorf("start command error: %v", err)
        return err
    }
    pid := int32(cmd.Process.Pid)
    for {
        line, err := output.ReadString('\n')
        if line != "" {
            srv.Send(&pb.RunResponse{Line: []string{line}, Pid: pid, Starttime: ts})
        }
        if err == io.EOF {
            break
        } else if err != nil {
            return err
        }
    }
    cmd.Wait()
    srv.Send(&pb.RunResponse{Finished: true, Success: cmd.ProcessState.Success()})
    return nil
}

func (pws *ProcWatcherService) scan() error {
    files, err := ioutil.ReadDir("/proc")
    if err != nil {
        log.Errorf("scan process error: %v", err)
        return err
    }
    for _, p := range pws.procs {
        p.Reset()
    }
    for _, f := range files {
        if !f.IsDir() {
            continue
        }
        pid, err := strconv.Atoi(f.Name())
        if err != nil || pid <= 1 {
            continue
        }
        b, u, ctime, err := ParseCmdline(pid)
        if err != nil {
            log.Errorf("parse process %d error: %v", pid, err)
            return err
        }
        for _, p := range pws.procs {
            cmdline := string(b)
            if p.reg == nil {
                p.reg, err = regexp.Compile(p.ProcessPattern)
                if err != nil {
                    return err
                }
            }
            if !p.reg.MatchString(cmdline) {
                continue
            }
            if p.ProcessUser != "" && p.ProcessUser != u.Name {
                continue
            }
            if err := p.process.Add(*u, ctime, pid); err != nil {
                return err
            }
        }
    }
    pws.lastScanTime = time.Now().Unix()
    return nil
}
